%load_ext autoreload
%autoreload 2

# deal with matplotlib plots
%matplotlib inline

# display all outputs in a cell
get_ipython().ast_node_interactivity = 'all'

from IPython.display import display, HTML

# misc
import sys
import os
import numpy as np
import random
from tqdm import tqdm
import pandas as pd

# io
import h5py

# plotting
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn; seaborn.set_style('white')
import matplotlib.ticker as ticker
import plotnine as pn
import logomaker
import tangermeme
from tangermeme.plot import plot_logo
from tangermeme.utils import reverse_complement
import matplotlib.animation as animation

sys.path.append("..")
from tangermeme_utils.utils import plot_motif
from chrombpnet_utils.modisco import trim_cwm

# editable text in PDFs
# https://stackoverflow.com/a/54111532
plt.rcParams['pdf.fonttype'] = 42
plt.rcParams['ps.fonttype'] = 42

INPUTLEN = 2114
OUTPUTLEN = 1000
SHIFT_REL_INPUT = np.int64((2114 - 1000) / 2)
SHIFT_REL_INPUT

NARROWPEAK_SCHEMA = ["chr", "start", "end", "1", "2", "3", "4", "5", "6", "summit"]

LOGO_ALPHABET = 'ACGT'

# finemo colorscheme
LOGO_COLORS = {"A": '#109648', "C": '#255C99', "G": '#F7B32B', "T": '#D62839'}

# logomaker colorscheme
LOGO_COLORS2= {
        'G': [1, .65, 0],
        'T': [1, 0, 0],
        'C': [0, 0, 1],
        'A': [0, .5, 0]
    }

CHAR_IGNORE = ['QWERYUIOPSDFHJKLZXVBNM']

with open("../../DURGA_DIRS.txt", 'r') as f:
	proj_in = f.readline().strip()
	proj_out = f.readline().strip()

with open("../../AK_PROJ_DIR.txt", 'r') as f:
    kundaje_dir = f.readline().strip()

cluster_meta = pd.read_csv(os.path.join(proj_out, "03-chrombpnet/01-models/qc/chrombpnet_models_keep2.tsv"), sep="\t", header=None)

# set column names
cluster_meta.columns = ["cluster", "folds", "cluster_ID"]

cluster_meta.head()

cluster_meta.shape

cluster = "Heart_c0"
path_modisco = proj_out + f"03-chrombpnet/01-models/modisco/bias_Heart_c0_thresh0.4/{cluster}/counts_modisco_output.h5"
modisco = h5py.File(path_modisco, 'r')
modisco.keys()

cwm = modisco['pos_patterns']['pattern_0']['contrib_scores'][()]
trimmed_cwm, _, _ = trim_cwm(cwm)
trimmed_cwm.shape
plot_motif(trimmed_cwm.T)
trimmed_cwm.sum()
trimmed_cwm.sum()/trimmed_cwm.shape[0]

cwm = modisco['pos_patterns']['pattern_10']['contrib_scores'][()]
trimmed_cwm, _, _ = trim_cwm(cwm)
trimmed_cwm.shape
plot_motif(trimmed_cwm.T)
trimmed_cwm.sum()
trimmed_cwm.sum()/trimmed_cwm.shape[0]

cwm = modisco['neg_patterns']['pattern_0']['contrib_scores'][()]
trimmed_cwm, _, _ = trim_cwm(cwm)
trimmed_cwm.shape
plot_motif(trimmed_cwm.T)
trimmed_cwm.sum()
trimmed_cwm.sum()/trimmed_cwm.shape[0]

# collect a dict of dfs
cwm_means = []

for cluster in cluster_meta['cluster'].values:
	
	print(cluster)
	path_modisco = proj_out + f"03-chrombpnet/01-models/modisco/bias_Heart_c0_thresh0.4/{cluster}/counts_modisco_output.h5"
	modisco = h5py.File(path_modisco, 'r')

	# collect a dict of dfs
	for grp in modisco.keys():

		for motif in modisco[grp].keys():

			# calculate sum of contribs and mean of contribs, across trimmed CWMs
			cwm = modisco[grp][motif]['contrib_scores'][()]
			trimmed_cwm, _, _ = trim_cwm(modisco[grp][motif]['contrib_scores'][()])
			cwm_sum = trimmed_cwm.sum()
			cwm_mean = cwm_sum/trimmed_cwm.shape[0]
			abs_cwm_sum = np.abs(trimmed_cwm).sum()
			abs_cwm_mean = abs_cwm_sum/trimmed_cwm.shape[0]

			cwm_means.append([cluster, grp, motif, cwm_sum, abs_cwm_sum, cwm_mean, abs_cwm_mean])

# convert list to df
cwm_means_df = pd.DataFrame(cwm_means, columns=["cluster", "pattern_class", "motif", "cwm_sum", "abs_cwm_sum", "cwm_mean", "abs_cwm_mean"])

# concatenate two columns
cwm_means_df['component_pattern'] = cwm_means_df['cluster'] + "__" + cwm_means_df['pattern_class'] + "." + cwm_means_df['motif']

cwm_means_df.head()

# write out
cwm_means_df.to_csv(os.path.join(proj_out, "03-chrombpnet/02-compendium/modisco_compiled/cwm_sums_df.tsv"), sep="\t", index=False)